import numpy as np
import calfem.geometry as cfg
import calfem.mesh as cfm
import calfem.vis as cfv
import os
import io
import sys
import time
import cProfile
import pstats
import matplotlib.pyplot as plt
from toolkits import gauss_line_int, gauss_tri_int, gauss_tri_weight, gauss_tri_point
from collections.abc import Iterable
from collections import deque
from copy import deepcopy
from queue import PriorityQueue

class HiddenPrints:
    def __enter__(self):
        self._original_stdout = sys.stdout
        sys.stdout = open(os.devnull, 'w')

    def __exit__(self, exc_type, exc_val, exc_tb):
        sys.stdout.close()
        sys.stdout = self._original_stdout

class Poisson2D:
    """docstring for Poisson
    
    3-Node triangular element FEM/EFEM model for 2D Poisson/Laplace equation. 

    """
    def __init__(self, geometry, DirMarker, NeuMarker, holeMarker, BCfunc, Source=None, direction=None):
        self.Geometry = geometry
        self.Node = None
        self.Edge = None
        self.Element = None
        self.DirMarker = DirMarker # list of markers for Dirichlet's boundary
        self.Dirichlet = None    # (i,j): marker
        self.NeuMarker = NeuMarker # list of markers for Neumann's boundary
        self.Neumann = None     # (i,j): marker
        self.NeuCluster = None
        self.holeMarker = holeMarker  # list of marker lists for each hole 
        self.hole = None      # (i,j): holeid 
        self.holeEdge = None  # holeid : [(i,j),(m,n),...]
        self.holeElement = None
        self.direction = direction
        self.crossEdge = None
        self.crossElement = None
        self.BCfunc = BCfunc
        self.Source = Source
        self.numHole = len(self.holeMarker)
        self.numElement = None
        self.numNode = None
        self.numDOF = None
        self.area = None
        # EFEM matrix vector
        self.FU = None
        self.dU = None
        self.F = None
        self.d = None
        self.PhiFe = None
        self.PhiLe = None
        self.s = None
        self.qh = None
        # FEM matrix vector
        self.KU = None
        self.fU = None
        self.K = None
        self.f = None
        self.PhiKe = None
        self.uh = None
        self.qh_fem = None
        # errors.
        self.flux_re_fem = None
        self.flux_re_efem = None
        self.elem_cre = None
        self.cre = None
        # private
        self._mesh = None
        self._edof = None
        self._BdElements = None
      
    def mesh(self, display=False, size=0.1):
        # generate mesh using GMSH.
        self._mesh = cfm.GmshMesh(self.Geometry, return_boundary_elements=True)
        self._mesh.elType = 2
        self._mesh.dofsPerNode = 1
        ##########################################
        #                                        #
        #          Basic Data Structure          #
        #                                        #
        # Node: list of coordinates (x, y)       #
        # Element: list of nodes (i,j,k) CCW.    #
        # Edge: dict of edges (i,j):n pair       #
        #                                        #
        ##########################################
        with HiddenPrints():
            self.Node, self._edof, _, _, _, self._BdElements = self._mesh.create()
        self.Element = self._edof - 1
        self.Edge = dict()
        for elemid, nodes in enumerate(self.Element):
            i,j,k = nodes
            self.Edge[(i, j)] = elemid
            self.Edge[(j, k)] = elemid
            self.Edge[(k, i)] = elemid

        self.numElement = len(self.Element)
        self.numNode = len(self.Node)
        self.numDOF = self.numNode + self.numHole

        ###################################################################
        #                                                                 #
        #           Geometry and Relationship Data Structure              #
        #                                                                 #
        # wholeElement: set of whole element ids.                         #
        # neighbor: dict of element's neighbors. elemnid: [elmenids] pair #
        # length: dict of edge's length  (i,j): length pair               #
        # area: area of each element   elemid: area pair                  #
        #                                                                 #
        ###################################################################
        self.wholeElement = set([x for x in range(self.numElement)])
        self.neighbor = {}
        self.length = {}
        self.area = np.zeros(self.numElement)
        
        for elemid, nodes in enumerate(self.Element):
            i,j,k = nodes
            x1, y1 = self.Node[i]
            x2, y2 = self.Node[j]
            x3, y3 = self.Node[k]

            self.neighbor[elemid] = set()
            for edge in [(j,i), (k,j), (i,k)]:
                if edge in self.Edge:
                    self.neighbor[elemid].add(self.Edge[edge])

            self.length[(i, j)] = np.linalg.norm(self.Node[i]-self.Node[j])
            self.length[(j, k)] = np.linalg.norm(self.Node[j]-self.Node[k])
            self.length[(k, i)] = np.linalg.norm(self.Node[k]-self.Node[i])            
            self.area[elemid] = 1/2*np.linalg.det(np.array([[x1, y1, 1],
                                                            [x2, y2, 1],
                                                            [x3, y3, 1]]))        

        ###############################################
        #                                             #
        #                 Marker Bank                 #
        #              (for applying BC)              #
        #                                             #
        # DirEdge: dict of marker, marker:(i,j)       #
        # NeuEdge: dict of marker, marker:(i,j)       #
        # holeEdge: dict of holeid, holeid:[(i,j),..] #
        #                                             #
        ###############################################
        self.hole = {}      # (i,j): holeid 
        self.holeEdge = {}  # holeid : [(i,j),(m,n),...]
        self.holeElement = {} # holeid: hole element set()

        for holeid, hole in enumerate(self.holeMarker):
            edges = []
            elements = set()
            for marker in hole:
                for e in self._BdElements[marker]:
                    edge = tuple(np.array(e['node-number-list'])-1)
                    # Force to CCW.
                    if edge not in self.Edge:
                        edge = edge[::-1]

                    self.hole[edge] = holeid

                    edges.append(edge)
                    elements.add(self.Edge[edge])
            self.holeEdge[holeid] = np.array(edges)
            self.holeElement[holeid] = elements

        ##############################################
        #                                            #
        #                Boundary sets               #
        #          (for applying constraint)         #
        #                                            #
        # Boundary: Boundary edge set.(i,j):marker   #
        # Dirichlet: Dirichlet edge set.(i,j):marker #
        # Neumann: Neumann edge set. (i,j):marker    #
        # hole: hole edge dict. (i,j):holeid         #  
        #                                            #
        ##############################################
        self.Dirichlet = {}
        self.Neumann = {}
        self.Boundary = {}
        self.boundaryElement = set() # boundary element set 

        # To do: should I cross to the externel edge? 
        # self.externelElement = set()

        for marker in self.DirMarker:
            for e in self._BdElements[marker]:
                edge = tuple(np.array(e['node-number-list'])-1)
                # Force to CCW.
                if edge not in self.Edge:
                    edge = edge[::-1]
                self.Dirichlet[edge] = marker
                self.Boundary[edge] = marker
                self.boundaryElement.add(self.Edge[edge])

        for marker in self.NeuMarker:
            for e in self._BdElements[marker]:
                edge = tuple(np.array(e['node-number-list'])-1)
                # Force to CCW.
                if edge not in self.Edge:
                    edge = edge[::-1]
                self.Neumann[edge] = marker
                self.Boundary[edge] = marker
                self.boundaryElement.add(self.Edge[edge])


        if len(self.Neumann) != 0:
            self.NeuCluster = self._cluster_edge(self.Neumann)
        ########################################
        #                                      #
        #     Crossing elements and edges      #
        #                                      #
        # crossEdge: (i,j):(n,eta)^T           #
        # crossElement: elemid: (rk_id, ae)^T  #
        #                          (tensor)    #
        ########################################
        self.crossEdge = dict()
        self.crossElement = dict()

        if len(self.holeMarker) != 0: 
            # self._connect_domain()
            self._one_piece()

        # Display mesh
        if display:
            self._check_ccw(size)
            cfv.draw_mesh(coords=self.Node, edof=self._edof, dofs_per_node=self._mesh.dofs_per_node,
                          el_type=self._mesh.el_type,filled=True,title="Mesh")
            cfv.showAndWait()

    def _cluster_edge(self, edges):
        # Attention to deepcopy!
        edges = deepcopy(edges)
        clusters = []

        e, _ = edges.popitem()
        clusters.append(deque([e]))
        
        while len(edges) != 0:
            quit = False
            for i, j in edges:
                for c in clusters:
                    if i == c[-1][-1]:
                        c.append((i, j))
                        edges.pop((i, j), None)
                        quit = True
                        break
                    if j == c[0][0]:
                        c.appendleft((i, j))
                        edges.pop((i, j), None)
                        quit = True
                        break
                if quit:
                    break
            else:
                if len(edges) != 0:
                    e, _ = edges.popitem()
                    clusters.append(deque([e]))
                    
        # convert to list
        clusters = [list(q) for q in clusters]
        return clusters 

    def _check_ccw(self, size=0.1):
        triangles = self.Element
        x = self.Node[:,0]
        y = self.Node[:,1]

        fig, axs = plt.subplots(1,2,dpi=300)
        
        axs[0].triplot(x,y,triangles,linewidth=0.1)
        axs[0].set_aspect('equal','box')
        axs[0].axis('off')

        if len(self.Neumann) != 0:
            for edge in self.Neumann:
                start = self.Node[edge[0]]
                end = self.Node[edge[1]]
                sx = start[0]
                sy = start[1]
                dx = end[0] - sx 
                dy = end[1] - sy
                axs[0].arrow(sx, sy, dx, dy, width=size/1000, head_width=size/10, length_includes_head=True, color='blue')
            
        if len(self.Dirichlet) != 0:
            for edge in self.Dirichlet:
                start = self.Node[edge[0]]
                end = self.Node[edge[1]]
                sx = start[0]
                sy = start[1]
                dx = end[0] - sx 
                dy = end[1] - sy
                axs[0].arrow(sx, sy, dx, dy, width=size/1000, head_width=size/10, length_includes_head=True, color='red')
          
        axs[1].triplot(x,y,triangles,linewidth=0.1)
        axs[1].set_aspect('equal','box')
        axs[1].axis('off')

        if len(self.Neumann) !=0:

            eff_dofs = []
            red_dofs = []

            for cluster in self.NeuCluster:
                ###############################################
                eff_dofs.append([cluster[0][0]])
                _red = []
                for edge in cluster:
                    # to avoid error occured in closed-edge.
                    if edge[1] != eff_dofs[-1][0]:
                        _red.append(edge[1])
                red_dofs.append(_red)
                ###############################################
                color = np.random.rand(3,)
                for edge in cluster:
                    start = self.Node[edge[0]]
                    end = self.Node[edge[1]]
                    sx = start[0]
                    sy = start[1]
                    dx = end[0] - sx 
                    dy = end[1] - sy
                    axs[1].arrow(sx, sy, dx, dy, width=size/1000, head_width=size/10, length_includes_head=True, color=color)
            
            all_red_dofs = [y for x in red_dofs for y in x]
            eff_dofs = [ dof for dofs in eff_dofs for dof in dofs]
            color = np.random.rand(3,)
            for dof in eff_dofs:
                axs[1].scatter(*self.Node[dof], color = color, s=size*100)

            color = np.random.rand(3,)
            for dof in all_red_dofs:
                axs[1].scatter(*self.Node[dof], color = color, s=size*50)


        if len(self.crossElement) != 0:
            for ei in self.crossElement:
                x = self.Node[self.Element[ei]]
                t = plt.Polygon(x)
                plt.gca().add_patch(t)

        plt.show()
    
    def _shortest_path(self, holeid):
        '''
        Dijkstra's algorithm for finding minimum number of elements connecting hole and boundary. 
        '''

        externelElement = self.boundaryElement - self.holeElement[holeid]

        Q = PriorityQueue()
        dist = {}
        prev = {}
        visited = set()

        for i in self.wholeElement:
            dist[i] = np.inf 

        for j in self.holeElement[holeid]:
            dist[j] = 0
            Q.put((dist[j], j))
        # print(externelElement)
        # print(self.boundaryElement)
        # print(self.holeElement[holeid])
        # print(dist, prev)
        flag = True
        final = None
        while not Q.empty() and flag:
            current = Q.get()[1]
            visited.add(current)

            # print('current :', current)
            # print('neighbor:', *self.neighbor[current])
            # print('distance:', *[dist[k] for k in self.neighbor[current]])
            for k in self.neighbor[current]:
                if k in externelElement:
                    dist[k] = dist[current] + 1
                    prev[k] = current
                    flag = False
                    final = k
                    # print('connected!', k)
                    break

                if k not in visited:
                    nd = dist[current] + 1
                    if dist[k] > nd:
                        dist[k] = nd
                        prev[k] = current
                        Q.put((dist[k], k))
            # print('distance:', *[dist[k] for k in self.neighbor[current]])

        path = []
        while True:
            if final in prev:
                path.append(final)
                final = prev[final]
            else:
                path.append(final)
                break
        path.reverse()
        return path

    def _one_piece(self):
        #################################
        #                               #
        #  Crossing elements and edges  #
        #                               #
        # crossEdge: (i,j):(n,eta)^T    #
        #                               #
        #################################        
        for holeid in self.holeEdge:
            path = self._shortest_path(holeid)
            # to do : no path

            ########################
            #                      #
            #      crossEdge       #
            # inlet edge eta = 1   #
            #                      #
            ########################
            i,j,k = self.Element[path[0]]
            candidates = [(i,j),(j,k),(k,i)]
            inlet = [x for x in candidates if x in self.holeEdge[holeid]][0]
            self.crossEdge[inlet] = np.array([[holeid+self.numNode], [1]])

            ###########################
            #                         #
            #       crossElement      #
            # modified shape function #
            #                         #
            ###########################
            n = len(path)

            for pid in range(n-1):
                elema = path[pid]
                elemb = path[pid+1]
                i,j,k = self.Element[elema]
                u,v,w = self.Element[elemb]
                nodea = [i,j,k]
                edgea = set([(i,j),(j,k),(k,i)])
                edgeb = set([(v,u),(w,v),(u,w)])
                outlet = edgea & edgeb
                outlet = outlet.pop()
                ae = np.zeros((3,1))
                ae[nodea.index(inlet[0])]= 1/self.length[inlet]
                ae[nodea.index(outlet[0])]= -1/self.length[outlet]
                ae = np.insert(ae, 0, holeid+self.numNode, axis=0)

                if elema in self.crossElement:
                    self.crossElement[elema] = np.c_[self.crossElement[elema], ae]
                else:
                    self.crossElement[elema] = ae
                inlet = outlet[::-1]

            # last element (also for n=1)
            last_elem = path[-1]
            i,j,k = self.Element[last_elem]
            last_nodes = [i,j,k]
            candidates = set([(i,j),(j,k),(k,i)])
            candidates.remove(inlet)
            outlet = [x for x in candidates if x in self.Boundary][0]

            ae = np.zeros((3,1))
            ae[last_nodes.index(inlet[0])]= 1/self.length[inlet]
            ae[last_nodes.index(outlet[0])]= -1/self.length[outlet]
            ae = np.insert(ae, 0, holeid+self.numNode, axis=0)
            if last_elem in self.crossElement:
                self.crossElement[last_elem] = np.c_[self.crossElement[last_elem], ae]
            else:
                self.crossElement[last_elem] = ae

            ########################
            #                      #
            # outlet edge eta = -1 #
            #                      #
            ########################
            if outlet in self.crossEdge:
                self.crossEdge[outlet] = np.c_[self.crossEdge[outlet], np.array([[holeid+self.numNode], [-1]])]
            else:
                self.crossEdge[outlet] = np.array([[holeid+self.numNode], [-1]])

    def __connect_domain(self):
        #################################
        #                               #
        #  Crossing elements and edges  #
        #                               #
        # crossEdge: (i,j):(n,eta)^T    #
        #                               #
        #################################

        if self.direction is None:
            directions = np.array([[0, -1]])
        else:
            directions = np.array(self.direction)

        for n in self.holeEdge:
            print('{} hole,'.format(n))
            direction = directions[n]
            edge = tuple(self.holeEdge[n][0])
            self.crossEdge[edge] = np.array([[n+self.numNode], [1]])

            _connected = False

            while not _connected:
                # print('current:', edge, self.Node[list(edge)][0], self.Node[list(edge)][1])
                ei = self.Edge[edge]
                
                i, j, k = self.Element[ei]
                order = [i,j,k]  
                
                local_edges = set([(i,j), (j,k), (k,i)])
                local_edges.remove(edge)
                e1 = np.array(local_edges.pop())
                e2 = np.array(local_edges.pop())
                
                e1n = self.Node[e1[1]] - self.Node[e1[0]]
                e1n = e1n/self.length[tuple(e1)]
                
                e2n = self.Node[e2[1]] - self.Node[e2[0]]
                e2n = e2n/self.length[tuple(e2)]
                
                d1 = np.abs(np.dot(direction, e1n))
                d2 = np.abs(np.dot(direction, e2n))
                ae = np.zeros((3,1))

                if d1 < d2:
                    L1 = self.length[edge]
                    L2 = self.length[tuple(e1)]
                    ae[order.index(edge[0])]= 1/L1
                    ae[order.index(e1[0])]= -1/L2
                    edge = tuple(e1)
                else:
                    L1 = self.length[edge]
                    L2 = self.length[tuple(e2)]
                    ae[order.index(edge[0])]= 1/L1
                    ae[order.index(e2[0])]= -1/L2
                    edge = tuple(e2)

                ae = np.insert(ae, 0, n+self.numNode, axis=0)

                if ei in self.crossElement:
                    self.crossElement[ei] = np.c_[self.crossElement[ei], ae]
                else:
                    self.crossElement[ei] = ae

                if edge in self.Boundary:
                    if edge in self.hole:
                        m = self.hole[edge]
                        if m == n:
                            # TODO: reject.
                            pass
                        else:
                            if edge in self.crossEdge:
                                self.crossEdge[edge] = np.c_[self.crossEdge[edge], np.array([[n+self.numNode], [-1]])]
                            else:
                                self.crossEdge[edge] = np.array([[n+self.numNode], [-1]])
                    else:
                        if edge in self.crossEdge:
                            self.crossEdge[edge] = np.c_[self.crossEdge[edge], np.array([[n+self.numNode], [-1]])]
                        else:
                            self.crossEdge[edge] = np.array([[n+self.numNode], [-1]])

                    _connected = True
                else:
                    edge = edge[::-1]

    def __prepare_q0(self, degree=2.5):
        # TODO: different Gaussian quadrature dgeree 

        # (k,) vector.
        self.Delta = np.repeat(self.area, 3)
        self.Wk = np.repeat(gauss_tri_weight[degree], self.numElement)
        self.b = np.zeros(self.numElement*3)
        self.xi = np.zeros(self.numElement*3)
        self.eta = np.zeros(self.numElement*3)
        # (j,) vector
        self.Wj = gauss_tri_weight[2]
        # (m,j) tensor
        self.alpha = np.zeros((self.numElement,3))
        self.beta = np.zeros((self.numElement,3))

        for i, e in enumerate(self.Element):
            x1, y1 = self.Node[e][0]
            x2, y2 = self.Node[e][1]
            x3, y3 = self.Node[e][2]

            shape = np.array([[x1-x3,y1-y3],
                              [x2-x3,y2-y3],
                              [x3, y3]])

            xi_eta_one = np.c_[gauss_tri_point[degree], np.ones(3)]
            alpha_beta_one = np.c_[gauss_tri_point[2], np.ones(3)]

            X = xi_eta_one @ shape
            Y = alpha_beta_one @ shape

            self.xi[i*3: (i+1)*3] = X[:,0]
            self.eta[i*3: (i+1)*3] = X[:,1]
            self.b[i*3: (i+1)*3] = self.Source(X[:,0], X[:,1])

            self.alpha[i] = Y[:,0]
            self.beta[i] = Y[:,1]

        # q_part is (k,) vector.
        self.q_part = self.Wk*self.b*self.Delta

    def _prepare_q0(self, degree=2.5):
        # TODO: different Gaussian quadrature dgeree 

        # # (k,) vector.
        # self.Delta = np.repeat(self.area, 3)
        # self.Wk = np.repeat(gauss_tri_weight[degree], self.numElement)

        # (n,) vector
        self.Delta = self.area.copy()
        # (i,) vector
        self.Wi = gauss_tri_weight[degree]
        # (n,i) tensor
        self.b = np.zeros((self.numElement,3))
        self.xi = np.zeros((self.numElement,3))
        self.eta = np.zeros((self.numElement,3))
        # (j,) vector
        self.Wj = gauss_tri_weight[2]
        # (m,j) tensor
        self.alpha = np.zeros((self.numElement,3))
        self.beta = np.zeros((self.numElement,3))

        for i, e in enumerate(self.Element):
            x1, y1 = self.Node[e][0]
            x2, y2 = self.Node[e][1]
            x3, y3 = self.Node[e][2]

            shape = np.array([[x1-x3,y1-y3],
                              [x2-x3,y2-y3],
                              [x3, y3]])

            xi_eta_one = np.c_[gauss_tri_point[degree], np.ones(3)]
            alpha_beta_one = np.c_[gauss_tri_point[2], np.ones(3)]

            X = xi_eta_one @ shape
            Y = alpha_beta_one @ shape

            self.xi[i] = X[:,0]
            self.eta[i] = X[:,1]
            self.b[i] = self.Source(X[:,0], X[:,1])
            self.alpha[i] = Y[:,0]
            self.beta[i] = Y[:,1]

        # q_part is (n,i ) tensor.
        self.q_part = (self.Wi*self.b)*self.Delta[...,np.newaxis]
        # self.q_part = self.Wk*self.b*self.Delta

    def _q0_x(self, x, y):
        if not isinstance(x, Iterable):
            # if x is scalar.
            # (n,i) tensor
            G = -1/(2*np.pi)*(x-self.xi)/((x-self.xi)**2+(y-self.eta)**2)
            # return inner of two (k,) vector
            return np.einsum('ni,ni->', G, self.q_part)
        else: 
            # if x, y are of (s,) shape.
            #  (s, ) vector to (s, n, i) tensor, than use broadcasting, (s, n, i) - (n, i)
            x = x[..., np.newaxis, np.newaxis]
            y = y[..., np.newaxis, np.newaxis]

            #  (s, n, i) tensor
            G = -1/(2*np.pi)*(x-self.xi)/((x-self.xi)**2+(y-self.eta)**2)
            # return (s,n,i),(n,i)->(s,)
            return np.einsum('sni,ni->s', G, self.q_part)

    def _q0_y(self, x, y):
        if not isinstance(x, Iterable):
            # if x is scalar.
            # (k,) vector
            G = -1/(2*np.pi)*(y-self.eta)/((x-self.xi)**2+(y-self.eta)**2)
            # return inner of two (k,) vector
            return np.einsum('ni,ni->', G, self.q_part)
        else: 
            # if x, y are of (s,) shape.
            #  (s, ) vector to (s, n, i) tensor, than use broadcasting, (s, n, i) - (n, i)
            x = x[..., np.newaxis, np.newaxis]
            y = y[..., np.newaxis, np.newaxis]

            #  (s, n, i) tensor
            G = -1/(2*np.pi)*(y-self.eta)/((x-self.xi)**2+(y-self.eta)**2)
            # return (s,n,i),(n,i)->(s,)
            return np.einsum('sni,ni->s', G, self.q_part)

    def _u0(self, x, y):
        if not isinstance(x, Iterable):
            # if x is scalar.
            # (n,i) tensor
            G = -1/(2*np.pi)*np.log(np.sqrt((x-self.xi)**2+(y-self.eta)**2))
            # return inner of two (k,) vector
            return np.einsum('ni,ni->', G, self.q_part)
        else: 
            # if x, y are of (s,) shape.
            #  (s, ) vector to (s, n, i) tensor, than use broadcasting, (s, n, i) - (n, i)
            x = x[..., np.newaxis, np.newaxis]
            y = y[..., np.newaxis, np.newaxis]

            #  (s, n, i) tensor
            G = -1/(2*np.pi)*np.log(np.sqrt((x-self.xi)**2+(y-self.eta)**2))
            # return (s,n,i),(n,i)->(s,)
            return np.einsum('sni,ni->s', G, self.q_part)        

    def __q0_x(self, x, y):
        if not isinstance(x, Iterable):
            # if x is scalar.
            # (k,) vector
            G = -1/(2*np.pi)*(x-self.xi)/((x-self.xi)**2+(y-self.eta)**2)
            # return inner of two (k,) vector
            return np.einsum('ni,ni->', G, self.q_part)
            return np.inner(G, self.q_part)
        else: 
            # if x, y are of (s,) shape.
            k = len(self.xi)
            #  (s, ) vector to (s, k) tensor, than use broadcasting, (s, k) - (k)
            x = np.tile(x[...,np.newaxis], (1,k))
            y = np.tile(y[...,np.newaxis], (1,k))

            #  (s, k) tensor
            G = -1/(2*np.pi)*(x-self.xi)/((x-self.xi)**2+(y-self.eta)**2)
            # return (s,k),(k,)->(s,)
            return np.einsum('sk,k->s', G, self.q_part)

    def __q0_y(self, x, y):
        if not isinstance(x, Iterable):
            # if x is scalar.
            # (k,) vector
            G = -1/(2*np.pi)*(y-self.eta)/((x-self.xi)**2+(y-self.eta)**2)
            # return inner of two (k,) vector
            return np.inner(G, self.q_part)
        else: 
            # if x, y are of (s,) shape.
            k = len(self.xi)
            #  (s, ) vector to (s, k) tensor, than use broadcasting, (s, k) - (k)
            x = np.tile(x[...,np.newaxis], (1,k))
            y = np.tile(y[...,np.newaxis], (1,k))

            #  (s, k) tensor
            G = -1/(2*np.pi)*(y-self.eta)/((x-self.xi)**2+(y-self.eta)**2)
            # return (s,k),(k,)->(s,)
            return np.einsum('sk,k->s', G, self.q_part)

    def _integral_q0(self):
        #  (m, j) tensor to (m, j, n, i) tensor, than use broadcasting, (m, j, n, i) - (n, i)
        Alpha = self.alpha[...,np.newaxis, np.newaxis]
        Beta = self.beta[...,np.newaxis, np.newaxis]

        # (m, j, n, i) tensor
        Gx_mjni = -1/(2*np.pi)*(Alpha-self.xi)/((Alpha-self.xi)**2+(Beta-self.eta)**2)
        Gy_mjni = -1/(2*np.pi)*(Beta-self.eta)/((Alpha-self.xi)**2+(Beta-self.eta)**2)

        # (m,j,n,i), (n,i) -> (m,j)
        Gx_mj = np.einsum('mjni,ni-> mj', Gx_mjni, self.q_part)
        Gy_mj = np.einsum('mjni,ni-> mj', Gy_mjni, self.q_part)
    
        # (m,j), (j,) -> (m,)
        # (m,)*(m,) -> (m,)

        self.int_q0 = np.zeros((self.numElement, 2))
        self.int_q0[:,0] = np.einsum('mj,j->m', Gx_mj, self.Wj) * self.area
        self.int_q0[:,1] = np.einsum('mj,j->m', Gy_mj, self.Wj) * self.area

    def _assemble_efem(self):
        ########################################
        ########################################
        ##                                    ##
        ##  unconstrained flexibility matrix  ##
        ##                                    ##
        ########################################
        ########################################
        self.FU = np.zeros((self.numDOF, self.numDOF))
        self.PhiFe = np.zeros((self.numElement, 2, 3))
        # specific for crossing element.
        self.PhiLe = dict() 

        for elemid, nodes in enumerate(self.Element):
            i,j,k = nodes
            x1, y1 = self.Node[i]
            x2, y2 = self.Node[j]
            x3, y3 = self.Node[k]
            
            self.PhiFe[elemid] = np.array([[x3-x2, x1-x3, x2-x1],
                                           [y3-y2, y1-y3, y2-y1]])/(2*self.area[elemid])    

            if elemid in self.crossElement:
                L1 = self.length[(i,j)]
                L2 = self.length[(j,k)]
                left = np.linalg.inv([[(y2-y1)/L1, (x1-x2)/L1],
                                      [(y3-y2)/L2, (x2-x3)/L2]])

                self.PhiLe[elemid] = np.c_[left, np.zeros((2,1))]

                PhiFe = np.c_[self.PhiFe[elemid], self.PhiLe[elemid] @ self.crossElement[elemid][1::,:]]
                Fe = self.area[elemid] * PhiFe.T @ PhiFe
                
                ey = np.r_[nodes, self.crossElement[elemid][0,:]].astype(int)

                self.FU[np.ix_(ey,ey)] = self.FU[np.ix_(ey,ey)] + Fe
            else:
                Fe = self.area[elemid] * self.PhiFe[elemid].T @ self.PhiFe[elemid]
                
                self.FU[np.ix_(nodes,nodes)] = self.FU[np.ix_(nodes,nodes)] + Fe

        print('size: {} x {}'.format(*self.FU.shape))
        # print('rank: {}'.format(np.linalg.matrix_rank(self.FU)))

        ########################################
        ########################################
        ##                                    ##
        ##   unconstrained kinematic vector   ##
        ##                                    ##
        ########################################
        ########################################
        self.dU = np.zeros(self.numDOF)
        #========================#
        #   Dirichlet boundary   #
        #========================#
        for edge in self.Dirichlet:
            i,j = edge
            L = self.length[edge]
            marker = self.Dirichlet[edge]
            
            up = gauss_line_int(self.BCfunc[marker], *self.Node[i], *self.Node[j])
            self.dU[i] = self.dU[i] - up/L
            self.dU[j] = self.dU[j] + up/L
            # the effect cause by rk term. 
            if edge in self.crossEdge:
                self.dU[self.crossEdge[edge][0]] = self.dU[self.crossEdge[edge][0]] + self.crossEdge[edge][1]*up/L

        #============#
        #   Source   #
        #============#
        if self.Source is not None:
            for i, e in enumerate(self.Element):
                if i in self.crossElement:
                    PhiFe = np.c_[self.PhiFe[i], self.PhiLe[i] @ self.crossElement[i][1::,:]]
                    dU_mod = PhiFe.T @ self.int_q0[i]
                    
                    ey = np.r_[e.copy(), self.crossElement[i][0,:]].astype(int)
                    self.dU[ey] = self.dU[ey] - dU_mod
                else:
                    dU_mod = self.PhiFe[i].T @ self.int_q0[i]
                    
                    self.dU[e] = self.dU[e] - dU_mod

        ## apply divergence theorem.
        # if self.Source is not None:
        #     for edge in self.Boundary:
        #         i,j = edge
        #         L = self.length[edge]
                
        #         up = gauss_line_int(self._u0, *self.Node[i], *self.Node[j])
        #         self.dU[i] = self.dU[i] + up/L
        #         self.dU[j] = self.dU[j] - up/L
        #         # the effect cause by rk term. 
        #         if edge in self.crossEdge:
        #             self.dU[self.crossEdge[edge][0]] = self.dU[self.crossEdge[edge][0]] - self.crossEdge[edge][1]*up/L



        # if self.Source is not None:
        #     for edge in self.Boundary:
        #         i, j = edge
        #         elemid = self.Edge[edge]

        #         elem = self.Element[elemid]

        #         x1, y1 = self.Node[i]
        #         x2, y2 = self.Node[j]
        #         L = self.length[edge]
        #         normal = np.array([[y2-y1],[x1-x2]])/L
                
        #         up = gauss_line_int(self._u0, x1, y1, x2, y2)

        #         if elemid in self.crossElement:
        #             PhiFe = np.c_[self.PhiFe[elemid], self.PhiLe[elemid] @ self.crossElement[elemid][1::,:]]
        #             dU_mod = PhiFe.T @ normal * up
        #             ey = np.r_[elem, self.crossElement[elemid][0,:]].astype(int)
        #         else:
        #             dU_mod = self.PhiFe[elemid].T @ normal * up
        #             ey = elem

        #         self.dU[ey] = self.dU[ey] - dU_mod.flatten()

    def _constraint_efem(self):
        #========================================#
        #  1.  Neumann boundary (without source) #
        #========================================#

        self.sN = np.zeros(self.numDOF)
        self.A = np.eye(self.numDOF)

        if len(self.Neumann) != 0:
            # todo : move to subfunction.
            eff_dofs = []
            red_dofs = []
            for clusters in self.NeuCluster:
                eff_dofs.append([clusters[0][0]])
                _red = []
                for edge in clusters:
                    # to avoid error occured in closed-edge.
                    if edge[1] != eff_dofs[-1][0]:
                        _red.append(edge[1])
                red_dofs.append(_red)


            for n, d in enumerate(eff_dofs):
                self.A[np.ix_(red_dofs[n], d)] = 1

            all_red_dofs = [y for x in red_dofs for y in x]

            # Neumann terms
            for cluster in self.NeuCluster:
                # Records
                rn = {}
                fl_bar = 0
                # # exclude the last edge in closed edge.
                # if cluster[0][0] == cluster[-1][-1]:
                #     cn = cluster[:-1]
                # else:
                #     cn = cluster

                cn = cluster
                for edge in cn:
                    marker = self.Neumann[edge]
                    i, j = edge
                    x1, y1 = self.Node[i]
                    x2, y2 = self.Node[j]
                    L = self.length[edge]

                    #============#
                    #   Source   #
                    #============#
                    if self.Source is not None:
                        neu_func = lambda x,y: self.BCfunc[marker](x,y) - (self._q0_x(x,y)*(y2-y1) + self._q0_y(x,y)*(x1-x2))/L
                    else:
                        neu_func = self.BCfunc[marker]
                    
                    fl_bar = fl_bar + gauss_line_int(neu_func, *self.Node[i], *self.Node[j])
                    self.sN[j] = fl_bar

                    if edge in self.crossEdge:
                        dofs, etas = self.crossEdge[edge]
                        for dof, eta in zip(dofs, etas):
                            rn[int(dof)] = -eta

                    # if rn is not empty, then always set rn
                    for dof in rn:
                        self.A[j, dof] = rn[dof]


            self.A = np.delete(self.A, all_red_dofs, axis=1)

        self.F = self.A.T @ self.FU @ self.A
        self.d = self.A.T @ (self.dU - self.FU @ self.sN)

        # print(np.allclose(self.F, self.FU), np.allclose(self.d, self.dU))
        #==============================================#
        #  2.  Uniqueness of nodal variables (s0 = 0)  #
        #==============================================#
        # self.F[:,0] # new 
        self.F[0,:] = 0
        self.F[0,0] = 1
        self.d[0] = 0

        #=====================================#
        #  3.  Neumann hole (rk is enforced)  #
        #=====================================#
        if len(self.Neumann) != 0:
            for cluster in self.NeuCluster:
                # 封闭的一定是孔吗？
                if (cluster[0][0] == cluster[-1][-1]) and (cluster[0] in self.hole):
                    # print('neumann hole!')
                    holeid = self.hole[cluster[0]]
                    dof = holeid - self.numHole
                    # print('hole{}, dof:{}'.format(holeid, dof))
                    r = 0
                    for edge in cluster:
                        marker = self.Neumann[edge]
                        i, j = edge
                        x1,y1 = self.Node[i]
                        x2,y2 = self.Node[j]
                        L = self.length[edge]
                        #============#
                        #   Source   #
                        #============#
                        if self.Source is not None:
                            neu_func = lambda x,y: self.BCfunc[marker](x,y) - (self._q0_x(x,y)*(y2-y1) + self._q0_y(x,y)*(x1-x2))/L
                        else:
                            neu_func = self.BCfunc[marker]

                        r = r + gauss_line_int(neu_func, x1, y1, x2, y2)
                    # print(dof, r)
                    self.F[dof,:] = 0
                    self.F[dof,dof] = 1
                    self.d[dof] = r
                    # print(r)
        # print(np.linalg.matrix_rank(self.F))

    def _solve_efem(self):
        _s = np.linalg.inv(self.F) @ self.d
        self.s = self.A @ _s + self.sN

    def _construct_flux(self):
        self.qh = np.zeros((self.numElement, 2)) 

        for i, e in enumerate(self.Element):
            
            if i in self.crossElement:                
                PhiFe = np.c_[self.PhiFe[i], self.PhiLe[i] @ self.crossElement[i][1::,:]]

                ey = np.r_[e, self.crossElement[i][0,:]].astype(int)

                self.qh[i] = PhiFe @ self.s[ey]

            else:
                self.qh[i] = self.PhiFe[i] @ self.s[e]

            if self.Source is not None:
                x, y = np.mean(self.Node[e],axis=0)
                self.qh[i] = self.qh[i] + np.array([self._q0_x(x,y), self._q0_y(x,y)])

    def _display_efem(self):
        with HiddenPrints():
            cfv.figure()
            cfv.drawElementValues(
                ev=self.qh[:,0],
                coords=self.Node,
                edof=self._edof,
                dofs_per_node=self._mesh.dofsPerNode,
                el_type=self._mesh.elType,
                title="EFEM flux field",
                draw_elements=True
                    )
            cfv.showAndWait()
            cfv.figure()
            cfv.drawElementValues(
                ev=self.qh[:,1],
                coords=self.Node,
                edof=self._edof,
                dofs_per_node=self._mesh.dofsPerNode,
                el_type=self._mesh.elType,
                title="EFEM flux field",
                draw_elements=True
                    )
            cfv.showAndWait()

    def error_efem(self, ux_func, uy_func):
        self.flux_re_efem = np.zeros(self.numElement)

        err = 0
        for elemid, e in enumerate(self.Element):
            i,j,k = e
            q_err_func = lambda x,y: (ux_func(x,y)-self.qh[elemid][0])**2+(uy_func(x,y)-self.qh[elemid][1])**2
            
            elem_err = gauss_tri_int(q_err_func, *self.Node[i], *self.Node[j], *self.Node[k])

            err += elem_err

            self.flux_re_efem[elemid] = elem_err

        return np.sqrt(err)

    def efem(self, display=False, log=False):
        if self._mesh == None:
            print('Please use .meshing() method to create mesh at first!')
        else:
            if log:
                pr = cProfile.Profile()
                pr.enable()
            if self.Source is not None:
                self._prepare_q0()
                self._integral_q0()
            self._assemble_efem()
            self._constraint_efem()
            self._solve_efem()
            self._construct_flux()
            if log:
                pr.disable()
                s = io.StringIO()
                ps = pstats.Stats(pr, stream=s).sort_stats('tottime')
                ps.print_stats()
                with open('log.txt', 'a+') as f:
                    f.write(s.getvalue())
            if display:
                self._display_efem()

    def _assemble_fem(self):
        ####################################
        #                                  #
        #  unconstrained stiffness matrix  #
        #                                  #
        ####################################
        self.KU = np.zeros((self.numNode, self.numNode))
        self.PhiKe = []

        for elemid, nodes in enumerate(self.Element):
            x1, y1 = self.Node[nodes][0]
            x2, y2 = self.Node[nodes][1]
            x3, y3 = self.Node[nodes][2]
            
            a1 = x2*y3 - x3*y2
            a2 = x3*y1 - x1*y3
            a3 = x1*y2 - x2*y1
            b1 = y2 - y3
            b2 = y3 - y1
            b3 = y1 - y2 
            c1 = x3 - x2
            c2 = x1 - x3 
            c3 = x2 - x1 

            self.PhiKe.append([lambda x,y : (a1 + b1*x + c1*y)/(2*self.area[elemid]),
                               lambda x,y : (a2 + b2*x + c2*y)/(2*self.area[elemid]),
                               lambda x,y : (a3 + b3*x + c3*y)/(2*self.area[elemid])])

            DerPhiKe = np.array([[b1, b2, b3],
                                 [c1, c2, c3]])/(2*self.area[elemid])
            

            Ke = self.area[elemid] * DerPhiKe.T @ DerPhiKe
                
            self.KU[np.ix_(nodes,nodes)] = self.KU[np.ix_(nodes,nodes)] + Ke

        print('size: {} x {}'.format(*self.KU.shape))
        # print('rank: {}'.format(np.linalg.matrix_rank(self.KU)))

        ######################################
        #                                    #
        #   unconstrained kinematic vector   #
        #                                    #
        ######################################
        self.fU = np.zeros(self.numNode)
        #======================#
        #   Neumann boundary   #
        #======================#
        for edge in self.Neumann:
            marker = self.Neumann[edge]
            elemid = self.Edge[edge]
            l,m,n = self.Element[elemid]
            edges = set([l,m,n])
            i, j = edge
            edges.remove(i)
            edges.remove(j)
            k = edges.pop()

            # f1 = lambda x,y: self.PhiKe[elemid][0](x,y)*self.BCfunc[marker](x,y)
            # f2 = lambda x,y: self.PhiKe[elemid][1](x,y)*self.BCfunc[marker](x,y)
            # f3 = lambda x,y: self.PhiKe[elemid][2](x,y)*self.BCfunc[marker](x,y)

            # # print(self.BCfunc[marker](*self.Node[i]), self.BCfunc[marker](*self.Node[j]))

            # fp1 = gauss_line_int(f1, *self.Node[i], *self.Node[j])
            # fp2 = gauss_line_int(f2, *self.Node[i], *self.Node[j])
            # fp3 = gauss_line_int(f3, *self.Node[i], *self.Node[j])
            # # print(l,m,n,i,j)

            # self.fU[l] = self.fU[l] + fp1
            # self.fU[m] = self.fU[m] + fp2
            # self.fU[n] = self.fU[n] + fp3
            g1 = self.BCfunc[marker](*self.Node[i])
            g2 = self.BCfunc[marker](*self.Node[j])
            g3 = 0
            g = np.array([g1,g2,g3])

            L = np.linalg.norm(self.Node[i]-self.Node[j])
            f = L*np.array([[1/3, 1/6, 0],
                            [1/6, 1/3, 0],
                            [  0,   0, 0]]) @ g

            self.fU[[i,j,k]] = self.fU[[i,j,k]] + f
        #============#  
        #   Source   #
        #============#
        if self.Source is not None:
            for elemid, tri in enumerate(self.Element):
                i,j,k = tri
                f1 = self.Source(*self.Node[i])
                f2 = self.Source(*self.Node[j])
                f3 = self.Source(*self.Node[k])
                f = np.array([f1,f2,f3])
                S = self.area[elemid]*np.array([[1/6,1/12,1/12],
                                                [1/12,1/6,1/12],
                                                [1/12,1/12,1/6]]) @ f
                self.fU[tri] = self.fU[tri] + S
                # f1 = lambda x,y: self.PhiKe[elemid][0](x,y)*self.Source(x,y)
                # f2 = lambda x,y: self.PhiKe[elemid][1](x,y)*self.Source(x,y)
                # f3 = lambda x,y: self.PhiKe[elemid][2](x,y)*self.Source(x,y)

                # s1 = gauss_tri_int(f1, *self.Node[i], *self.Node[j], *self.Node[k])
                # s2 = gauss_tri_int(f2, *self.Node[i], *self.Node[j], *self.Node[k])
                # s3 = gauss_tri_int(f3, *self.Node[i], *self.Node[j], *self.Node[k])

                # S = np.array([s1,s2,s3])
                # self.fU[tri] = self.fU[tri] + S

    def _constraint_fem(self):
        self.K = self.KU.copy()
        self.f = self.fU.copy()
  
        for edge in self.Dirichlet:
            marker = self.Dirichlet[edge]
            i, j = edge

            self.K[i,:] = 0
            self.K[i,i] = 1
            self.f[i] = self.BCfunc[marker](*self.Node[i])      

            self.K[j,:] = 0
            self.K[j,j] = 1
            self.f[j] = self.BCfunc[marker](*self.Node[j])

        if len(self.Dirichlet) == 0:
            self.K[0,:] = 0
            self.K[0,0] = 1
            self.f[0] = 1
        
        # print(np.linalg.matrix_rank(self.K))

    def _construct_flux_fem(self):
        self.qh_fem = np.zeros((self.numElement,2))

        for elemid, node in enumerate(self.Element):
            x1, y1 = self.Node[node][0]
            x2, y2 = self.Node[node][1]
            x3, y3 = self.Node[node][2]
            
            b1 = y2 - y3
            b2 = y3 - y1
            b3 = y1 - y2 
            c1 = x3 - x2
            c2 = x1 - x3 
            c3 = x2 - x1 

            DerPhiKe = np.array([[b1, b2, b3],
                                 [c1, c2, c3]])/(2*self.area[elemid])

            self.qh_fem[elemid] = DerPhiKe @ self.uh[node]

    def _solve_fem(self):

        self.uh = np.linalg.inv(self.K) @ self.f 

    def _display_fem(self):
        u = np.zeros(len(self.Element))

        for i, e in enumerate(self.Element):
            u[i] = np.mean(self.uh[e])
        
        with HiddenPrints():
            cfv.figure()
            cfv.drawElementValues(
                ev=u,
                coords=self.Node,
                edof=self._edof,
                dofs_per_node=self._mesh.dofsPerNode,
                el_type=self._mesh.elType,
                title="fem temperature field",
                draw_elements=True
                    )
            cfv.showAndWait()
            cfv.figure()
            cfv.drawElementValues(
                ev=self.qh_fem[:,0],
                coords=self.Node,
                edof=self._edof,
                dofs_per_node=self._mesh.dofsPerNode,
                el_type=self._mesh.elType,
                title="fem temperature field",
                draw_elements=True
                    )
            cfv.showAndWait()
            cfv.figure()
            cfv.drawElementValues(
                ev=self.qh_fem[:,1],
                coords=self.Node,
                edof=self._edof,
                dofs_per_node=self._mesh.dofsPerNode,
                el_type=self._mesh.elType,
                title="fem temperature field",
                draw_elements=True
                    )
            cfv.showAndWait()

    def error_fem(self, ux_func, uy_func):
        self.flux_re_fem = np.zeros(self.numElement)

        err = 0
        for elemid, e in enumerate(self.Element):
            x1, y1 = self.Node[e][0]
            x2, y2 = self.Node[e][1]
            x3, y3 = self.Node[e][2]
            
            b1 = y2 - y3
            b2 = y3 - y1
            b3 = y1 - y2 
            c1 = x3 - x2
            c2 = x1 - x3 
            c3 = x2 - x1 

            DerPhiKe = np.array([[b1, b2, b3],
                                 [c1, c2, c3]])/(2*self.area[elemid])

            qxh, qyh = DerPhiKe @ self.uh[e]

            i,j,k = e
            q_err_func = lambda x,y: (ux_func(x,y)-qxh)**2+(uy_func(x,y)-qyh)**2

            elem_err = gauss_tri_int(q_err_func, *self.Node[i], *self.Node[j], *self.Node[k])

            err += elem_err
            self.flux_re_fem[elemid] = elem_err

        return np.sqrt(err)

    def fem(self, display=False, log=False):
        if self._mesh == None:
            print('Please use .meshing() method to create mesh at first!')
        else:
            if log:
                pr = cProfile.Profile()
                pr.enable()
            self._assemble_fem()
            self._constraint_fem()
            self._solve_fem()
            self._construct_flux_fem()
            if log:
                pr.disable()
                s = io.StringIO()
                ps = pstats.Stats(pr, stream=s).sort_stats('tottime')
                ps.print_stats()
                with open('log.txt', 'a+') as f:
                    f.write(s.getvalue())
            if display:
                self._display_fem()

    def display_exact(self, func):
        q = np.zeros(len(self.Element))

        for i, e in enumerate(self.Element):
            x, y = np.mean(self.Node[e], axis=0)
            q[i] = func(x, y)
        with HiddenPrints():
            cfv.figure()
            cfv.drawElementValues(
                ev=q,
                coords=self.Node,
                edof=self._edof,
                dofs_per_node=self._mesh.dofsPerNode,
                el_type=self._mesh.elType,
                title="Exact field",
                draw_elements=True
                    )
            cfv.showAndWait()
            
    def calculate_cre(self):
        # self.elem_cre = np.zeros(self.numElement)
        # self.cre = 0

        # err = np.sum(np.power(self.qh-self.qh_fem, 2),axis=1) 

        # err = self.area*err

        # cre = np.sqrt(np.sum(err))
        

        # err = 0
        # for elemid, node in enumerate(self.Element):

        #     # print(self.qh_fem[elemid], self.qh[elemid])

        #     q_err_inner = np.inner(self.qh_fem[elemid] - self.qh[elemid], self.qh_fem[elemid] - self.qh[elemid])
            
        #     elem_err = self.area[elemid]*q_err_inner

        #     # self.elem_cre[elemid] = elem_err

        #     err += elem_err

        # self.cre = np.sqrt(err)

        # return self.cre

        err = 0
        for elemid, nodes in enumerate(self.Element):
            i,j,k = nodes        
            norm = (self.qh[elemid][0]-self.qh_fem[elemid][0])**2+(self.qh[elemid][1]-self.qh_fem[elemid][1])**2
            q_err_func = lambda x,y: norm*np.ones_like(x)

            elem_err = gauss_tri_int(q_err_func, *self.Node[i], *self.Node[j], *self.Node[k])

            err += elem_err

        return np.sqrt(err)
    # def calculate_cre(self):
    #     self.elem_cre = np.zeros(self.numElement)
    #     self.cre = 0

    #     err = 0
    #     for elemid, node in enumerate(self.Element):
    #         x1, y1 = self.Node[node][0]
    #         x2, y2 = self.Node[node][1]
    #         x3, y3 = self.Node[node][2]
            
    #         b1 = y2 - y3
    #         b2 = y3 - y1
    #         b3 = y1 - y2 
    #         c1 = x3 - x2
    #         c2 = x1 - x3 
    #         c3 = x2 - x1 

    #         DerPhiKe = np.array([[b1, b2, b3],
    #                              [c1, c2, c3]])/(2*self.area[elemid])

    #         grad_uh = DerPhiKe @ self.uh[node]

    #         q_err_inner = np.inner(grad_uh - self.qh[elemid], grad_uh - self.qh[elemid])
            
    #         elem_err = self.area[elemid]*q_err_inner

    #         self.elem_cre[elemid] = elem_err

    #         err += elem_err

    #     self.cre = np.sqrt(err)

    #     return self.cre 


class ConvergenceRate:
    """docstring for ConvergenceRate"""
    def __init__(self, model, geometry_func):
        self.model = model
        self.geometry_func = geometry_func
        self.hs_efem = None
        self.hs_fem = None
        self.dofs_efem = None
        self.dofs_fem = None
        self.errs_efem = None
        self.errs_fem = None
        self.conds_fem = None
        self.conds_efem = None
        self.dof_slope_fem = 0
        self.size_slope_fem = 0
        self.dof_slope_efem = 0
        self.size_slope_efem =  0
        self.dof_slope_cre = 0
        self.size_slope_cre = 0


    def convergence_efem(self, exact_ux, exact_uy, start=2, end=-1, num=6, display=False):
        self.hs_efem = np.logspace(start, end, num)
        self.errs_efem = []
        for h in self.hs_efem:
            self.model.Geometry = self.geometry_func(h)
            self.model.mesh()
            self.model.efem()
            self.errs_efem.append(self.model.error_efem(exact_ux, exact_uy))
        if display:
            self.plot_converage_rate(self.hs_efem, self.errs_efem)


    def convergence_fem(self, exact_ux, exact_uy, start=2, end=-1, num=6, display=False):
        self.hs_fem = np.logspace(start, end, num)
        self.errs_fem = []
        for h in self.hs_fem:
            self.model.Geometry = self.geometry_func(h)
            self.model.mesh()
            self.model.fem()
            self.errs_fem.append(self.model.error_fem(exact_ux, exact_uy))
        if display:
            self.plot_converage_rate(self.hs_fem, self.errs_fem)

    def plot_converage_rate(self, h, e):
        import matplotlib.pyplot as plt
        plt.figure()
        h = [1/x for x in h]
        plt.plot(h, e)
        plt.scatter(h, e, marker='o')
        plt.xscale("log")
        plt.yscale("log")
        plt.grid()
        plt.show()

    def convergence_all(self, exact_ux, exact_uy, start=2, end=-1, num=6, display=False, condition=False):
        self.hs_fem = np.logspace(start, end, num)
        self.dofs_fem = []
        self.errs_fem = []
        self.hs_efem = np.logspace(start, end, num)
        self.dofs_efem = []
        self.errs_efem = []
        self.conds_fem = []
        self.conds_efem = []
        self.cres = []
        self.dof_slope_fem = 0
        self.size_slope_fem = 0
        self.dof_slope_efem = 0
        self.size_slope_efem =  0
        self.dof_slope_cre = 0
        self.size_slope_cre = 0


        for h in self.hs_fem:
            self.model.Geometry = self.geometry_func(h)
            self.model.mesh(display=display)
            self.model.fem()
            self.model.efem()
            self.dofs_fem.append(self.model.numDOF)
            self.dofs_efem.append(self.model.numDOF)
            self.errs_fem.append(self.model.error_fem(exact_ux, exact_uy))
            self.errs_efem.append(self.model.error_efem(exact_ux, exact_uy))
            self.cres.append(self.model.calculate_cre())
            if condition:
                self.conds_fem.append(np.linalg.cond(self.model.K))
                self.conds_efem.append(np.linalg.cond(self.model.F))
            print('Calculation complete for mesh size:{:.4f}'.format(h))

        A = np.c_[np.log(self.dofs_fem), np.ones_like(self.dofs_fem)]
        y = np.log(self.errs_fem)
        k, b = np.linalg.inv(A.T@A) @ (A.T@y)
        self.dof_slope_fem = k

        A = np.c_[np.log(self.hs_fem), np.ones_like(self.hs_fem)]
        y = np.log(self.errs_fem)
        k, b = np.linalg.inv(A.T@A) @ (A.T@y)
        self.size_slope_fem = k

        A = np.c_[np.log(self.dofs_efem), np.ones_like(self.dofs_efem)]
        y = np.log(self.errs_efem)
        k, b = np.linalg.inv(A.T@A) @ (A.T@y)
        self.dof_slope_efem = k

        A = np.c_[np.log(self.hs_efem), np.ones_like(self.hs_efem)]
        y = np.log(self.errs_efem)
        k, b = np.linalg.inv(A.T@A) @ (A.T@y)
        self.size_slope_efem =  k     

        A = np.c_[np.log(self.dofs_efem), np.ones_like(self.dofs_efem)]
        y = np.log(self.cres)
        k, b = np.linalg.inv(A.T@A) @ (A.T@y)
        self.dof_slope_cre = k

        A = np.c_[np.log(self.hs_efem), np.ones_like(self.hs_efem)]
        y = np.log(self.cres)
        k, b = np.linalg.inv(A.T@A) @ (A.T@y)
        self.size_slope_cre =  k   


        if condition:
            print("Convergence rate calculation is complete!")
            print("Condition number:")
            print("Size \t|  FEM \t| EFEM  ")
            for h, fem, efem in zip(self.hs_fem, self.conds_fem, self.conds_efem):
                print('{:.2f} \t| {:.2e} \t| {:.2e} '.format(h, fem, efem))



    def plot_multiple_convergence_rate(self, which='size', file_name=None, font_size=18, fem_kwargs={}, efem_kwargs={}, fem_aux_kwargs={}, efem_aux_kwargs={}):
        if (self.hs_efem is None) or (self.hs_fem is None):
            print("Please call convergence_all() method first!")
            return

        hs_group = [self.hs_fem, self.hs_efem]
        dofs_group = [self.dofs_fem, self.dofs_efem]
        errs_group = [self.errs_fem, self.errs_efem]
        dofs_slope_group = [self.dof_slope_fem, self.dof_slope_efem]
        size_slope_group = [self.size_slope_fem, self.size_slope_efem]

        kwargs_group = [fem_kwargs, efem_kwargs] 
        aux_kwargs_group = [fem_aux_kwargs, efem_aux_kwargs] 


        import matplotlib.pyplot as plt
        import matplotlib

        matplotlib.rc('text', usetex=True)
        matplotlib.rc('text.latex', preamble=r'\usepackage{amsmath}')

        fig, ax = plt.subplots(figsize=(4,6), dpi=300)

        if which == "size":
            for h, e, s, kwargs in zip(hs_group, errs_group, size_slope_group, kwargs_group):
                kwargs['label'] = kwargs['label'] + '(Slope: {:.2f})'.format(s)
                ax.plot(h, e, **kwargs)

            # for h, aux_kwargs in zip(hs_group, aux_kwargs_group):
            #     aux = [aux_kwargs["C"]*x**-1 for x in h]
            #     del aux_kwargs["C"]
            #     ax.plot(h, aux, **aux_kwargs)
            # ax.set_xlabel(r"Inverse of mesh size, $\frac{1}{h}$")
        elif which == "dofs":
            for d, e, s, kwargs in zip(dofs_group, errs_group, dofs_slope_group, kwargs_group):
                kwargs['label'] = kwargs['label'].format(s)
                ax.plot(d, e, **kwargs)

            # for d, aux_kwargs in zip(dofs_group, aux_kwargs_group):
            #     aux = [aux_kwargs["C"]*x**-1 for x in d]
            #     del aux_kwargs["C"]
            #     ax.plot(d, aux, **aux_kwargs)
            # ax.set_xlabel(r"Degrees of freedom, $n$")
        else:
            print("The parameter 'which' should be 'size' or 'dofs'!")
            return    

        ax.tick_params(axis="both", which="both", direction="in")    
        ax.tick_params(axis='both', which='major', labelsize=font_size)
        ax.tick_params(axis='both', which='major', length=8)
        ax.tick_params(axis='both', which='major', width=1)
        ax.tick_params(axis='both', which='minor', length=4)
        ax.tick_params(axis='both', which='minor', width=1)

        # ax.tick_params(axis = 'both', which = 'minor', labelsize = 24)

        ax.set_xscale("log")
        ax.set_yscale("log")
        # ax.get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter())
        
        # ax.set_ylabel(r"Error in norm, $e$")
        # ax.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
        # ax.xaxis.set_major_formatter(matplotlib.ticker.FormatStrFormatter('%d'))
        ax.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
        ax.xaxis.set_minor_formatter(matplotlib.ticker.ScalarFormatter())
        ax.xaxis.set_minor_formatter(matplotlib.ticker.LogFormatter(labelOnlyBase=True, minor_thresholds=(0.5,0.2)))

        # ax.grid()
        plt.setp(ax.spines.values(), linewidth=2)
        plt.legend(fontsize='x-large')
        if file_name is None:
            plt.savefig("multiple_convergence.png")
        else:
            plt.savefig(file_name)
        plt.show()
        

